{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import warnings\n",
    "from QKDNetwork import QKDNetwork\n",
    "from compare import Compare\n",
    "from docplex.mp.model import Model\n",
    "from add_ps import AddPS\n",
    "import networkx as nx\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "import threading\n",
    "\n",
    "matplotlib.rcParams['svg.fonttype'] = 'none'\n",
    "warnings.filterwarnings('ignore', category=RuntimeWarning)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[28.984375, 19.4375, 8.90625, nan]\n",
      "[28.703125, 30.875, 11.859375, nan]\n"
     ]
    }
   ],
   "source": [
    "def get_sps_add_ps_num(net: QKDNetwork, init_ps: list):\n",
    "    sps_count = 0\n",
    "    for node in net.G.nodes:\n",
    "        if(node in init_ps):\n",
    "            net.G.nodes[node][\"transmitter\"] = 1\n",
    "    for sd in net.sd_list:\n",
    "        path = nx.shortest_path(net.G, sd[0], sd[1])\n",
    "        for path_index, path_node in enumerate(path):\n",
    "            if(path_index % 2 == 0 and net.G.nodes[path_node][\"transmitter\"] == 0):\n",
    "                net.G.nodes[path_node][\"transmitter\"] = 1\n",
    "                sps_count += 1\n",
    "    return sps_count\n",
    "\n",
    "def get_ras_add_ps_num(net: QKDNetwork, init_ps: list):\n",
    "    ras_count = 0\n",
    "    for node in net.G.nodes:\n",
    "        if(node in init_ps):\n",
    "            net.G.nodes[node][\"transmitter\"] = 1\n",
    "    for edge in net.G.edges:\n",
    "        if(net.G.nodes[edge[0]][\"transmitter\"] == 0 and net.G.nodes[edge[1]][\"transmitter\"] == 0):\n",
    "            ras_count += 1\n",
    "            if(net.G.degree(edge[0]) > net.G.degree(edge[1])):\n",
    "                net.G.nodes[edge[0]][\"transmitter\"] = 1\n",
    "            else:\n",
    "                net.G.nodes[edge[1]][\"transmitter\"] = 1\n",
    "    return ras_count\n",
    "\n",
    "def clear_all_placed_ps(net: QKDNetwork):\n",
    "    for node in net.G.nodes:\n",
    "        if(net.G.nodes[node][\"transmitter\"] == 1):\n",
    "            net.G.nodes[node][\"transmitter\"] = 0\n",
    "\n",
    "# 节点个数： 20-210          默认：60      迭代：20 + cnt * 10\n",
    "# 源目对数： 10-200          默认：50      迭代：10 + cnt * 10\n",
    "# 阿尔法值： 0.35-1.3        默认：0.2     迭代：0.2 + cnt * 0.05\n",
    "# 初始光源： 5-43            默认：5       迭代：5 + cnt * 2\n",
    "default_node_num = 60\n",
    "default_sd_num = 50\n",
    "default_alpha = 0.2\n",
    "default_init_ps_num = 5\n",
    "\n",
    "all_fig_data = []\n",
    "remove_ilp = True\n",
    "remove_ilp_rounding = False\n",
    "\n",
    "def simulation_run(fig_idx, fig_index, fig_data_single):\n",
    "    q = None\n",
    "    init_ps = []\n",
    "    if(fig_idx == 0):\n",
    "        q = QKDNetwork(num_node=20 + fig_index * 10, sd_num=default_sd_num, alpha=default_alpha, no_add_ps=True)\n",
    "    elif(fig_idx == 1):\n",
    "        q = QKDNetwork(num_node=default_node_num, sd_num=10 + fig_index * 10, alpha=default_alpha, no_add_ps=True)\n",
    "    elif(fig_idx == 2):\n",
    "        q = QKDNetwork(num_node=default_node_num, sd_num=default_sd_num, alpha=0.2 + fig_index * 0.05, no_add_ps=True)\n",
    "    elif(fig_idx == 3):\n",
    "        q = QKDNetwork(num_node=default_node_num, sd_num=default_sd_num, alpha=default_alpha, no_add_ps=True)\n",
    "    init_ps = random.sample(list(range(q.num_nodes)), default_init_ps_num + (fig_idx == 3) * fig_index * 1)\n",
    "    clear_all_placed_ps(q)\n",
    "    RAS_data = get_ras_add_ps_num(q, init_ps)\n",
    "    if(RAS_data != 0):\n",
    "        fig_data_single[0].append(RAS_data)\n",
    "    clear_all_placed_ps(q)\n",
    "    SPS_data = get_sps_add_ps_num(q, init_ps)\n",
    "    if(SPS_data != 0):\n",
    "        fig_data_single[1].append(SPS_data)\n",
    "    clear_all_placed_ps(q)\n",
    "    ret = AddPS(q).run_expr(init_ps, remove_ilp, remove_ilp_rounding)\n",
    "    if(ret[0] != 0):\n",
    "        fig_data_single[2].append(ret[0])\n",
    "    if(ret[1] != 0):\n",
    "        fig_data_single[3].append(ret[1])\n",
    "\n",
    "def thread_function(fig_idx, fig_index, fig_data_single):\n",
    "    simulation_run(fig_idx, fig_index, fig_data_single)\n",
    "\n",
    "for fig_idx in range(1, 4):\n",
    "    fig_data = [[], [], [], []]\n",
    "    for fig_index in range(10):  # 每张图多少x轴点\n",
    "        threads = []\n",
    "        fig_data_single = [[], [], [], []]\n",
    "        for _ in range(64): \n",
    "            t = threading.Thread(target=thread_function, args=(fig_idx, fig_index, fig_data_single))\n",
    "            threads.append(t)\n",
    "            t.start()\n",
    "        for t in threads:\n",
    "            t.join()  \n",
    "        fig_data[0].append(np.mean(fig_data_single[0]))\n",
    "        fig_data[1].append(np.mean(fig_data_single[1]))\n",
    "        fig_data[2].append(np.mean(fig_data_single[2]))\n",
    "        fig_data[3].append(np.mean(fig_data_single[3]))\n",
    "        \n",
    "        print([np.mean(fig_data_single[0]), np.mean(fig_data_single[1]), np.mean(fig_data_single[2]), np.mean(fig_data_single[3])])\n",
    "        \n",
    "    print(fig_data)\n",
    "    all_fig_data.append(fig_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(all_fig_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 节点个数： 20-210          默认：60      迭代：20 + cnt * 10\n",
    "# 源目对数： 10-200          默认：50      迭代：10 + cnt * 10\n",
    "# 阿尔法值： 0.35-1.3        默认：0.2     迭代：0.2 + cnt * 0.05\n",
    "# 初始光源： 5-43            默认：5       迭代：5 + cnt * 2\n",
    "\n",
    "for fig_index, data in enumerate(all_fig_data):\n",
    "    x_arr = []\n",
    "    if(fig_index == 0):\n",
    "        x_arr = list(range(20, len(data[0]) * 10 + 20, 10))\n",
    "    if(fig_index == 1):\n",
    "        x_arr = list(range(10, len(data[0]) * 10 + 10, 10))\n",
    "    if(fig_index == 2):\n",
    "        x_arr = list(np.arange(0.2, len(data[0]) * 0.05 + 0.2, 0.05))\n",
    "    if(fig_index == 3):\n",
    "        x_arr = list(range(5, len(data[0]) * 2 + 5, 2))\n",
    "    \n",
    "    plt.figure(figsize=(8, 6))\n",
    "    plt.plot(x_arr, data[0], label='RAS', marker='p')\n",
    "    plt.plot(x_arr, data[1], label='SPS', marker='o')\n",
    "    plt.plot(x_arr, data[2], label='RPSP-APSP', marker='s')\n",
    "    plt.plot(x_arr, data[3], label='RPSP-ILP', marker='^')\n",
    "\n",
    "    plt.xlabel(['Number of Nodes','Number of QKD Request','Network Connectivity','Number of PS nodes at the beginning'][fig_index])\n",
    "    plt.ylabel('Number of PSes to be added')\n",
    "    plt.legend()\n",
    "    plt.grid(True)\n",
    "    # plt.savefig(f\"var_add_ps_num_{fig_index}.svg\", format=\"svg\")\n",
    "    plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "xs",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
